import re
import xml.etree.cElementTree as ET
import numpy
import pickle
import os


################################################################################

def readDataBlock(xmlnode):
    """ Turn any 'DataBlock' XML node into a numpy array of floats
    """
    vmin = float(xmlnode.get('min'))
    vmax = float(xmlnode.get('max'))
    string = xmlnode.text
    string = re.sub("[\t\s\n]", "", string)
    data = numpy.asarray(
        bytearray.fromhex(string),
        dtype=float
    )
    return data * (vmax - vmin) / 255. + vmin


class Candidate(object):
    def __init__(self, fname):
        """ Build a new Candidate object from a PHCX file path.
        """
        xmlroot = ET.parse(fname).getroot()

        # Separate PDMP & FFT sections
        for section in xmlroot.findall('Section'):
            if 'pdmp' in section.get('name').lower():
                opt_section = section
            else:
                fft_section = section

        ### Sub-Integrations
        subintsNode = opt_section.find('SubIntegrations')
        nsubs = int(subintsNode.get('nSub'))
        nbins = int(subintsNode.get('nBins'))
        print("subint", nsubs, nbins)
        self.subints = readDataBlock(subintsNode).reshape(nsubs, nbins)

        ### Sub-Bands
        subbandsNode = opt_section.find('SubBands')
        nsubs = int(subbandsNode.get('nSub'))
        nbins = int(subbandsNode.get('nBins'))
        print("subband", nsubs, nbins)
        self.subbands = readDataBlock(subbandsNode).reshape(nsubs, nbins)



def extract_pulsars_features():

    # root_path = '/home/sunzy/workspace/data/MedlatTrainingData/'
    root_path = '/usr/data/MedlatTrainingData/'
    # candicate_path_base = '/home/sunzy/workspace/data/MedlatTrainingData/pulsars/'
    candicate_path_base = '/usr/data/MedlatTrainingData/pulsars/'
    phcx_files = os.listdir(candicate_path_base)
    print(phcx_files)
    pulsars_subints = {}
    pulsars_subbands = {}
    for phcx_file in phcx_files:
        id = phcx_file.split(".")[0]
        file_path = os.path.join(candicate_path_base, phcx_file)
        cand = Candidate(file_path)
        pulsars_subints[id] = cand.subints
        pulsars_subbands[id] = cand.subbands
    print(len(pulsars_subints), len(pulsars_subbands))
    output_path = os.path.join(root_path, "feature")
    # for py3
    # with open(os.path.join(output_path, "pulsars_subints"), "w") as f:
    # for py2
    with open(os.path.join(output_path, "pulsars_subints"), "wb") as f:
        pickle.dump(pulsars_subints, f)
    with open(os.path .join(output_path, "pulsars_subbands"), "wb") as f:
        pickle.dump(pulsars_subbands, f)

def extract_non_pulsars_features():
    # root_path = '/home/sunzy/workspace/data/MedlatTrainingData/'
    root_path = '/usr/data/MedlatTrainingData/'
    for i in range(1):# 9 for windows
        # dir_path = "RFI (%d)" % (i+1)
        dir_path = "RFI"
        candicate_path_base = os.path.join(root_path, dir_path)
        print(candicate_path_base)
        phcx_files = os.listdir(candicate_path_base)
        # print(phcx_files)

        dict_subints = {}
        dict_subbands = {}
        for phcx_file in phcx_files:
            id = phcx_file.split(".")[0]
            file_path = os.path.join(candicate_path_base, phcx_file)
            cand = Candidate(file_path)
            dict_subints[id] = cand.subints
            dict_subbands[id] = cand.subbands

        print(len(dict_subints), len(dict_subbands))
        output_path = os.path.join(root_path, "feature")
        # for py3
        # with open(os.path.join(output_path, "RFI_subints" ), "w") as f:
        # for py2
        with open(os.path.join(output_path, "RFI_subints"), "wb") as f:
            pickle.dump(dict_subints, f)
        with open(os.path.join(output_path, "RFI_subbands"), "wb") as f:
            pickle.dump(dict_subbands, f)

def test():
    """
    read the npy file
    :return:
    """
    # root_path = '/home/sunzy/workspace/data/MedlatTrainingData/'
    root_path = '/usr/data/MedlatTrainingData/'
    output_path = os.path.join(root_path, "feature")
    # for py3
    # with open(os.path.join(output_path, "RFI_subints"), "r") as f:
    # for py2
    with open(os.path.join(output_path, "RFI_subints"), "rb") as f:
        d = pickle.load(f)
    print(len(d))
    print(d.keys())

def main():
    extract_pulsars_features()
    extract_non_pulsars_features()
    test()


if __name__ == '__main__':
    main()
